#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#  Copyright (C) 2020 ShenZhen Powersmart Information Technology Co.,Ltd
#  All Rights Reserved.
#  本软件为深圳博沃智慧开发研制。未经本公司正式书面同意，其他任何个人、团体不得使用、
#  复制、修改或发布本软件.

# @Time : 2025/9/10 09:21
# @File : pollution_type_service.py
# @Project:
# @Author : wanghongtao
# @Software: PyCharm
import json
import re
from typing import Dict, Any, Optional, Tuple
from dataclasses import dataclass

from pyexpat.errors import messages
from sqlalchemy.orm import Session

from core import get_db
from services.ai_service_factory import ServiceType, AIService, ServiceConfig, ServiceRequest, ServiceResponse, \
    ServiceStatus
from services.visits_services.llm_analysis_services import LLMRunService
from services.pollution_type_service import pollution_type_service
from core.logging_config import get_logger

# 日志
logger = get_logger("pollution_type_service")

# 提示词模板
PROMPT_TEMPLATE = """# 角色
环保信访投诉线索打标签的专员

## 标签定义
{{label_define}}

{{label_example}}

## 信访投诉
{{xfnr}}
{{ts_dfnr}}

## 分析要求
请根据以下步骤完成分析任务：
1. **污染来源识别**：从信访投诉中提取出噪声发生的具体地点、行为、设备或人群。
2. **场景分类**：判断该噪声属于哪种场景（河道、湖泊、工厂、企业、建筑施工、道路运输、公共场所、家庭、学校、商业场所等）。
3. **标签匹配**：对照上述所有标签定义，选择最符合的标签，只识别最符合的一个。
4. **判断依据**：列出支持该标签选择的关键信息。
5. **置信度评估**：基于信息完整性和描述清晰度，给出 0~1 的置信度评分。

## 输出格式
1、根据场景和污染类型，分析[标签定义]中的定义，然后严格场景按照[标签定义]打标签，禁止编造新的标签，以JSON格式输出。
2、输出格式: ```json
{"污染类型":"标签定义的类型名称，只能取定义内的标签序号，例如标签1、标签2等",
"判断依据":"选择标签的原因，仔细分析不同标签的判断原因",
"置信度":"污染类型判断的置信度，判断范围：0~1"}```"""

# 配置定义
first_label_define = {
    "环境管理": ["环境管理", "建设项目审批"],
    "噪声污染": ["噪音污染"],
    "大气污染": ["大气污染"],
    "水污染": ["水污染"],
    "固体废物和土壤污染": ["固体污染", "危险化学品污染"],
    "辐射污染": ["电磁污染", "放射"],
    "其他污染": ["其他", "生态破坏_生态破坏", "海洋污染"],
}

hjgl_label_define = {
    "环境管理": {
        "环境监测": ["环境监测"],
        "环境执法": ["环境执法"],
        "其他环境管理": ["淘汰落后产能补偿", "生态示范和模范城区创建", "污染损害纠纷调解", "环境资质管理", "其他"]
    },
    "建设项目审批和排污许可": {
        "环境评价": ["环境评估", "未批先建", "建设项目环境影响评价审查信息公开"],
        "排污许可": [],
        "其他资质许可": [],
        "其他建设项目审批和排污许可": ["项目审批"]
    }
}


class PollutionTypeRecognitionService(LLMRunService):
    """
    污染类型识别服务

    继承自LLMRunService，利用大模型进行污染类型的逐级识别
    """

    def __init__(self, config: ServiceConfig = None):
        """初始化服务"""
        if config is None:
            config = ServiceConfig(
                service_type=ServiceType.POLLUTION_TYPE_RECOGNITION,
                name="污染类型识别服务",
                description="基于大模型的污染类型逐级识别服务",
                version="1.0.0"
            )
        super().__init__(config)
        self.config = config
        self.pollution_type_service = pollution_type_service
        self.first_label_define = first_label_define
        self.hjgl_label_define = hjgl_label_define
        self.prompt_template = PROMPT_TEMPLATE
        self.status = ServiceStatus.READY

    async def initialize(self) -> bool:
        """初始化服务"""
        try:
            logger.info(f"初始化污染类型识别服务")

            # 调用父类初始化
            parent_result = await super().initialize()
            if not parent_result:
                return False

            logger.info(f"污染类型识别服务初始化成功")
            self.status = ServiceStatus.READY
            return True

        except Exception as e:
            logger.error(f"污染类型识别服务初始化失败: {str(e)}")
            self.status = ServiceStatus.ERROR
            return False

    async def process(self, request: ServiceRequest) -> ServiceResponse:
        """处理污染类型识别请求"""
        start_time = self._get_current_time()

        try:
            db = next(get_db())
            data = request.data
            if not data or not data.get("xh"):
                error_msg = "缺少必要的数据字段: xh"
                logger.error(error_msg)
                return ServiceResponse(
                    request_id=request.request_id,
                    service_type=self.config.service_type,
                    success=False,
                    error_message=error_msg,
                    processing_time=self._get_current_time() - start_time
                )

            # 执行污染类型识别
            result = await self.analyze_pollution_type(data, db)
            new_result = data|result
            return ServiceResponse(
                request_id=request.request_id,
                service_type=self.config.service_type,
                success=True,
                data=new_result,
                processing_time=self._get_current_time() - start_time
            )

        except Exception as e:
            error_msg = f"处理污染类型识别请求失败: {str(e)}"
            logger.error(error_msg)
            return ServiceResponse(
                request_id=request.request_id,
                service_type=self.config.service_type,
                success=False,
                error_message=error_msg,
                processing_time=self._get_current_time() - start_time
            )

    async def analyze_pollution_type(self, data: Dict[str, Any], db: Session) -> Dict[str, Any]:
        """
        污染类型识别主流程
        """
        result = {
            "petition_record_id": data["xh"],
            "tslx_yj": "",
            "tslx_yjyj": "",
            "tslx_yj_zxd": "",
            "tslx_rj": "",
            "tslx_rjyj": "",
            "tslx_rj_zxd": "",
            "tslx_sj": "",
            "tslx_sjyj": "",
            "tslx_sj_zxd": ""
        }

        # 获取客户ID，默认使用"default"
        customer_id = data.get("customer_id", "default")

        # 1. 识别一级标签
        first_label = self._preprocess_first_label(data)
        if first_label:
            result["tslx_yj"] = first_label
            result["tslx_yjyj"] = f"基于分类内容字段快速识别：{data.get('flnr', '')}"
            result["tslx_yj_zxd"] = "1.0"
        else:
            # 使用大模型识别一级标签
            first_result = await self._recognize_by_llm(data, "first", db, customer_id=customer_id)
            result.update(first_result)

        # 2. 识别二级标签
        if result["tslx_yj"] == "环境管理":
            # 使用环境管理特殊处理
            second_label, third_label = self._preprocess_hjgl_label(data)
            if second_label:
                result["tslx_rj"] = second_label
                result["tslx_rjyj"] = f"基于内容分类字段识别：{data.get('内容分类', '')}"
                result["tslx_rj_zxd"] = "1.0"
                if third_label:
                    result["tslx_sj"] = third_label
                    result["tslx_sjyj"] = f"基于内容分类字段识别：{data.get('内容分类', '')}"
                    result["tslx_sj_zxd"] = "1.0"
        else:
            # 使用大模型识别二级标签
            second_result = await self._recognize_by_llm(data, "second", db, result["tslx_yj"], customer_id=customer_id)
            result.update(second_result)

            # 识别三级标签
            if result["tslx_rj"]:
                third_result = await self._recognize_by_llm(data, "third", db, result["tslx_rj"], customer_id=customer_id)
                result.update(third_result)

        return result

    def _preprocess_first_label(self, data: Dict[str, Any]) -> str:
        """预处理一级标签"""
        try:
            nrfl = data.get("flnr", "")
            for key, words in self.first_label_define.items():
                for word in words:
                    if word in nrfl:
                        return key
            return ""
        except Exception as e:
            logger.error(f"预处理一级标签失败: {str(e)}")
            return ""

    def _preprocess_hjgl_label(self, data: Dict[str, Any]) -> Tuple[str, str]:
        """预处理环境管理标签"""
        second_lable = ""
        thrid_label = ""
        try:
            nrfl = data.get("flnr", "")
            nrfl = str(nrfl)
            nrfl_list = nrfl.split("_")
            if len(nrfl_list) == 3:
                new_second_label = nrfl_list[1]
                new_thrid_lable = nrfl_list[2]

                # 处理环境管理类型
                if new_second_label == "环境管理":
                    second_labe_define = self.hjgl_label_define.get("环境管理")
                    for k, v in second_labe_define.items():
                        if new_thrid_lable in v:
                            second_lable = k
                            return second_lable, ""

                # 处理建设项目审批类型
                elif new_second_label == "建设项目审批":
                    second_labe_define = self.hjgl_label_define.get("建设项目审批和排污许可")
                    for k, v in second_labe_define.items():
                        if new_thrid_lable in v:
                            second_lable = "建设项目审批和排污许可"
                            thrid_label = k
                            return second_lable, thrid_label

        except Exception as e:
            logger.error(f"识别环境管理标签异常：{e}")
        return second_lable, thrid_label

    def _get_label_definitions(self, db: Session, customer_id: str, level: str, parent_label: str = None):
        """
        从数据库获取标签定义

        Args:
            customer_id: 客户ID
            level: 标签级别 (first/second/third)
            parent_label: 父级标签名称

        Returns:
            str: 格式化的标签定义字符串
        """
        label_define_dict = {}
        try:
            # 定义参数
            readly_define_dict = {
                "tag_category": "**定义**",
                "usage_scope": "**适用范围**",
                "exclusion_scope": "**排除范围**",
                "description": "**典型示例**"
            }
            # 根据级别查询不同层级的标签
            if level == "first":
                # 获取一级标签
                types = self.pollution_type_service.get_all_pollution_types(
                    customer_id, active_only=True, db=db
                )
                first_level_types = [t for t in types if t["level"] == 1]

                definitions = []
                for i, t in enumerate(first_level_types, 1):
                    definitions.append(f"### 标签{i}: {t['name']}")
                    label_define_dict[f"标签{i}"] = t["name"]
                    for key, value in readly_define_dict.items():
                        if t.get(key, None):
                            definitions.append(f"  {value}: {t[key]}")

                return "\n".join(definitions), label_define_dict

            elif level == "second" and parent_label:
                # 获取二级标签
                parent_type = self.pollution_type_service.get_pollution_type_by_name(
                    customer_id, parent_label, db
                )
                if parent_type:
                    types = self.pollution_type_service.get_all_pollution_types(
                        customer_id, active_only=True, db=db
                    )
                    second_level_types = [t for t in types if t["level"] == 2 and t["parent_code"] == parent_type.code]

                    definitions = []
                    for i, t in enumerate(second_level_types, 1):
                        definitions.append(f"### 标签{i}: {t['name']}")
                        label_define_dict[f"标签{i}"] = t["name"]
                        for key, value in readly_define_dict.items():
                            if t.get(key, None):
                                definitions.append(f"  {value}: {t[key]}")

                    return "\n".join(definitions), label_define_dict

            elif level == "third" and parent_label:
                # 获取三级标签
                parent_type = self.pollution_type_service.get_pollution_type_by_name(
                    customer_id, parent_label, db
                )
                if parent_type:
                    types = self.pollution_type_service.get_all_pollution_types(
                        customer_id, active_only=True, db=db
                    )
                    third_level_types = [t for t in types if t["level"] == 3 and t["parent_code"] == parent_type.code]

                    definitions = []
                    for i, t in enumerate(third_level_types, 1):
                        definitions.append(f"### 标签{i}: {t['name']}")
                        label_define_dict[f"标签{i}"] = t["name"]
                        for key, value in readly_define_dict.items():
                            if t.get(key, None):
                                definitions.append(f"  {value}: {t[key]}")

                    return "\n".join(definitions), label_define_dict

            return None, label_define_dict

        except Exception as e:
            logger.error(f"获取标签定义失败: {str(e)}")
            return None

    def _get_label_examples(self, db: Session, customer_id: str, level: str, parent_label: str = None) -> str:
        """
        从数据库获取标签样例

        Args:
            customer_id: 客户ID
            level: 标签级别
            parent_label: 父级标签名称

        Returns:
            str: 格式化的样例字符串
        """
        try:
            examples = []

            if level == "first":
                # 获取一级标签样例
                types = self.pollution_type_service.get_all_pollution_types(
                    customer_id, active_only=True, db=db
                )
                first_level_types = [t for t in types if t["level"] == 1]

                for t in first_level_types[:3]:  # 限制样例数量
                    type_examples = self.pollution_type_service.get_pollution_type_examples(
                        customer_id, t["id"], active_only=True, db=db
                    )
                    if type_examples:
                        examples.append(f"### {t['name']}样例:")
                        for ex in type_examples[:2]:  # 每个类型最多2个样例
                            examples.append(f"- {ex['example_content']}")

            elif level == "second" and parent_label:
                # 获取二级标签样例
                parent_type = self.pollution_type_service.get_pollution_type_by_name(
                    customer_id, parent_label, db
                )
                if parent_type:
                    types = self.pollution_type_service.get_all_pollution_types(
                        customer_id, active_only=True, db=db
                    )
                    second_level_types = [t for t in types if t["level"] == 2 and t["parent_code"] == parent_type.code]

                    for t in second_level_types[:3]:  # 限制样例数量
                        type_examples = self.pollution_type_service.get_pollution_type_examples(
                            customer_id, t["id"], active_only=True, db=db
                        )
                        if type_examples:
                            examples.append(f"### {t['name']}样例:")
                            for ex in type_examples[:2]:  # 每个类型最多2个样例
                                examples.append(f"- {ex['example_content']}")

            elif level == "third" and parent_label:
                # 获取三级标签样例
                parent_type = self.pollution_type_service.get_pollution_type_by_name(
                    customer_id, parent_label, db
                )
                if parent_type:
                    types = self.pollution_type_service.get_all_pollution_types(
                        customer_id, active_only=True, db=db
                    )
                    third_level_types = [t for t in types if t["level"] == 3 and t["parent_code"] == parent_type.code]

                    for t in third_level_types[:3]:  # 限制样例数量
                        type_examples = self.pollution_type_service.get_pollution_type_examples(
                            customer_id, t["id"], active_only=True, db=db
                        )
                        if type_examples:
                            examples.append(f"### {t['name']}样例:")
                            for ex in type_examples[:2]:  # 每个类型最多2个样例
                                examples.append(f"- {ex['example_content']}")

            return "\n".join(examples) if examples else ""

        except Exception as e:
            logger.error(f"获取标签样例失败: {str(e)}")
            return "样例获取失败"

    # 使用父类的LLM服务能力直接处理
    def _parse_llm_response(self, response_data: Dict[str, Any]) -> Dict[str, Any]:
        """解析抽取结果"""
        try:
            content = response_data.content
            content = content.strip("\n")
            # 尝试解析JSON
            if content.startswith("{") and content.endswith("}"):
                json_obj = json.loads(content)
                # 整理结果
                return json_obj
            else:
                json_obj = json.loads(content)
                # 如果不是JSON，返回原始内容
                return json_obj

        except Exception as e:
            logger.error(f"解析抽取结果失败: {str(e)}")
            return {"error": str(e)}

    async def _recognize_by_llm(self, data: Dict[str, Any], level: str, db: Session,
                               parent_label: str = None, customer_id: str = "default") -> Dict[str, Any]:
        """
        使用大模型进行标签识别

        Args:
            data: 输入数据
            level: 识别级别
            parent_label: 父级标签
            customer_id: 客户ID

        Returns:
            Dict[str, Any]: 识别结果
        """
        try:
            # 获取标签定义和样例
            label_define, label_define_dict = self._get_label_definitions(db, customer_id, level, parent_label)
            label_example = self._get_label_examples(db, customer_id, level, parent_label)
            if label_define:
                # 渲染提示词
                rendered_prompt = self._render_prompt(
                    self.prompt_template,
                    label_define=label_define,
                    label_example=label_example,
                    xfnr=data.get("xfnr", ""),
                    ts_dfnr=data.get("ts_dfnr", ""),
                )
                llm_params = {
                    "model": "glm-4-flash",
                    "max_tokens": 1024,
                    "temperature": 0.1,
                    "provider":"zhipu"
                }
                llm_response = await self._call_llm_service(llm_params, rendered_prompt)

                if llm_response.get("success"):
                    # 解析大模型响应
                    content = llm_response.get("response", "")
                    parsed_result = self._parse_llm_response(content)

                    # 将识别结果转换为设定的污染标签
                    if parsed_result.get("污染类型", ""):
                        if parent_label:
                            parsed_result["污染类型"] = label_define_dict.get(
                                parsed_result.get("污染类型"), f"其他{parent_label}")
                        else:
                            parsed_result["污染类型"] = label_define_dict.get(
                                parsed_result.get("污染类型"), f"其他{parent_label}")
                    if parsed_result.get("置信度"):
                        parsed_result["置信度"] = float(parsed_result["置信度"])*100
                    # 根据级别返回对应字段
                    if level == "first":
                        return {
                            "tslx_yj": parsed_result.get("污染类型", ""),
                            "tslx_yjyj": parsed_result.get("判断依据", ""),
                            "tslx_yj_zxd": str(parsed_result.get("置信度", 0))
                        }
                    elif level == "second":
                        return {
                            "tslx_rj": parsed_result.get("污染类型", ""),
                            "tslx_rjyj": parsed_result.get("判断依据", ""),
                            "tslx_rj_zxd": str(parsed_result.get("置信度", 0))
                        }
                    elif level == "third":
                        return {
                            "tslx_sj": parsed_result.get("污染类型", ""),
                            "tslx_sjyj": parsed_result.get("判断依据", ""),
                            "tslx_sj_zxd": str(parsed_result.get("置信度", 0))
                        }

                return {}
            else:
                logger.info(f"污染类型{parent_label}对应的子类定义为空！")
                return {}

        except Exception as e:
            logger.error(f"大模型识别失败: {str(e)}")
            return {}

    def _render_prompt(self, template: str, **kwargs) -> str:
        """渲染提示词模板"""
        try:
            result = template
            for key, value in kwargs.items():
                result = result.replace(f"{{{{{key}}}}}", str(value))
            return result
        except Exception as e:
            logger.error(f"渲染提示词模板失败: {str(e)}")
            return template

    def _get_current_time(self) -> float:
        """获取当前时间"""
        import time
        return time.time()

    async def get_capabilities(self) -> Dict[str, Any]:
        """
        获取服务能力信息

        Returns:
            Dict[str, Any]: 服务能力描述
        """
        return {
            "service_type": self.config.service_type.value,
            "service_name": self.config.name,
            "service_description": self.config.description,
            "version": self.config.version,
            "capabilities": {
                "pollution_type_recognition": True,
                "hierarchical_classification": True,
                "multi_level_analysis": ["first", "second", "third"],
                "llm_based_analysis": True,
                "preprocessing_support": True,
                "customer_isolation": True,
                "confidence_scoring": True,
                "template_based_prompting": True
            },
            "supported_features": {
                "first_level_preprocessing": "基于flnr字段的快速识别",
                "environmental_management_special_handling": "环境管理类型的特殊处理逻辑",
                "dynamic_label_definitions": "从数据库动态获取标签定义",
                "label_examples": "支持标签样例辅助识别",
                "multi_tenant_support": "支持多客户隔离"
            },
            "input_requirements": {
                "required_fields": ["xh", "xfnr"],
                "optional_fields": ["flnr", "内容分类", "customer_id"],
                "data_formats": ["dict", "json"]
            },
            "output_format": {
                "petition_record_id": "string",
                "tslx_yj": "string",  # 一级标签
                "tslx_yjyj": "string",  # 一级标签判断依据
                "tslx_yj_zxd": "string",  # 一级标签置信度
                "tslx_rj": "string",  # 二级标签
                "tslx_rjyj": "string",  # 二级标签判断依据
                "tslx_rj_zxd": "string",  # 二级标签置信度
                "tslx_sj": "string",  # 三级标签
                "tslx_sjyj": "string",  # 三级标签判断依据
                "tslx_sj_zxd": "string"  # 三级标签置信度
            }
        }

    async def health_check(self) -> Dict[str, Any]:
        """
        健康检查方法

        Returns:
            Dict[str, Any]: 健康检查结果
        """
        start_time = self._get_current_time()
        db = next(get_db())

        try:
            # 检查服务状态
            service_status = {
                "status": self.status.value if hasattr(self, 'status') else "unknown",
                "is_ready": hasattr(self, 'status') and self.status == ServiceStatus.READY
            }

            # 检查LLM服务管理器
            llm_status = {
                "available": hasattr(self, 'llm_service_manager') and self.llm_service_manager is not None,
                "model_loaded": False
            }

            if llm_status["available"]:
                try:
                    # 尝试获取模型列表
                    models = await self.llm_service_manager.list_models()
                    llm_status["model_loaded"] = len(models) > 0
                    llm_status["available_models"] = len(models)
                except Exception as e:
                    llm_status["error"] = str(e)
                    llm_status["available"] = False

            # 检查污染类型服务
            pollution_service_status = {
                "available": hasattr(self, 'pollution_type_service') and self.pollution_type_service is not None,
                "database_connected": False
            }

            if pollution_service_status["available"]:
                try:
                    # 测试数据库连接
                    customer_id = "default"
                    types = self.pollution_type_service.get_all_pollution_types(customer_id, active_only=True, db=db)
                    pollution_service_status["database_connected"] = True
                    pollution_service_status["available_types"] = len(types)
                except Exception as e:
                    pollution_service_status["error"] = str(e)
                    pollution_service_status["database_connected"] = False

            # 检查标签定义
            label_definitions_status = {
                "first_level_available": False,
                "examples_available": False
            }

            if pollution_service_status["database_connected"]:
                try:
                    customer_id = "default"
                    # 测试获取一级标签定义
                    first_level_defs, _ = self._get_label_definitions(db, customer_id, "first")
                    label_definitions_status["first_level_available"] = first_level_defs != "未找到相关标签定义"

                    # 测试获取标签样例
                    examples = self._get_label_examples(db, customer_id, "first")
                    label_definitions_status["examples_available"] = examples != "暂无样例"
                except Exception as e:
                    label_definitions_status["error"] = str(e)

            # 计算整体健康状态
            overall_healthy = (
                service_status["is_ready"] and
                llm_status["available"] and
                pollution_service_status["database_connected"] and
                label_definitions_status["first_level_available"]
            )

            health_result = {
                "service_name": self.config.name,
                "service_type": self.config.service_type.value,
                "overall_status": "healthy" if overall_healthy else "unhealthy",
                "timestamp": start_time,
                "check_duration_ms": round((self._get_current_time() - start_time) * 1000, 2),
                "service_status": service_status,
                "llm_service": llm_status,
                "pollution_type_service": pollution_service_status,
                "label_definitions": label_definitions_status,
                "checks_passed": 0,
                "total_checks": 4
            }

            # 计算通过的检查数量
            if service_status["is_ready"]:
                health_result["checks_passed"] += 1
            if llm_status["available"]:
                health_result["checks_passed"] += 1
            if pollution_service_status["database_connected"]:
                health_result["checks_passed"] += 1
            if label_definitions_status["first_level_available"]:
                health_result["checks_passed"] += 1

            return health_result

        except Exception as e:
            # 健康检查本身失败
            return {
                "service_name": self.config.name,
                "service_type": self.config.service_type.value,
                "overall_status": "error",
                "timestamp": start_time,
                "check_duration_ms": round((self._get_current_time() - start_time) * 1000, 2),
                "error": f"Health check failed: {str(e)}",
                "service_status": {"status": "error", "is_ready": False},
                "checks_passed": 0,
                "total_checks": 4
            }


# 全局服务实例
# pollution_type_recognition_service = PollutionTypeRecognitionService()